package blacklist

import (
	"context"
	"fmt"

	"dfx_cmdb/cmdb-agent/global"
	"dfx_cmdb/cmdb-agent/types"
	"dfx_cmdb/common"

	"gitee.com/daofeixian/dfx_golang_toolset/dfxnftables"
	"github.com/google/nftables"
	"github.com/zeromicro/go-zero/core/logc"
)

type HaproxyBlackList struct {
}

func NewHaproxyBlackList() *HaproxyBlackList {
	return &HaproxyBlackList{}
}

var (
	_       AgentBlacklistInter = (*HaproxyBlackList)(nil)
	nfMeate                     = &dfxnftables.NfMeta{
		AddressFamily: dfxnftables.NFT_ADDRESS_FAMILY_NET,
		TableName:     global.CMDB_BLACKLIST_HAPROXY_NAME,
	}

	nfTable = &dfxnftables.Nftables{
		NfMeta: *nfMeate,
	}
)

type HaproxyBlackListInfo struct {
	Port   string
	IpList []string
}

// cleanupExistingRules 清理已存在的黑名单规则和set
// 该函数会删除所有引用指定set的规则，然后删除set本身
func (h *HaproxyBlackList) cleanupExistingRules(ctx context.Context, ipListMap map[string]types.StringList, nfta *nftables.Table, nfChain *dfxnftables.NfChain) error {
	logc.Infof(ctx, "开始清理已存在的黑名单规则和set")

	// 遍历所有需要处理的 set
	for setName := range ipListMap {
		nfSet := &dfxnftables.NfSet{
			NfMeta:  *nfMeate,
			SetName: setName,
		}

		// 检查 set 是否存在
		exists, err := nfSet.CheckSet(ctx, global.Svc.Nft)
		if err != nil {
			logc.Errorf(ctx, "检查nfset %s 失败，错误信息: %v", setName, err)
			return err
		}

		if !exists {
			logc.Debugf(ctx, "set %s 不存在，跳过清理", setName)
			continue
		}

		logc.Debugf(ctx, "set %s 已存在，准备删除旧规则和set", setName)

		// 先删除引用该 set 的规则
		nfRule := &dfxnftables.NfChainRule{
			NfMeta:    *nfMeate,
			ChainName: "input",
			SrcSet:    setName,
		}

		nfchd, err := nfChain.QueryChainHandle(ctx, global.Svc.Nft, nfta)
		if err != nil {
			logc.Errorf(ctx, "查询nfchain %s 失败，错误信息: %v", nfMeate.TableName, err)
			return err
		}

		// 创建 nfchainrule - 匹配源IP在set中的流量并drop

		nfub, err := nfRule.CheckRule(ctx, global.Svc.Nft, nfta, nfchd)
		if err != nil {
			logc.Errorf(ctx, "查询nfchainrule %s 失败，错误信息: %v", nfMeate.TableName, err)
			return err
		}

		if nfub {
			logc.Infof(ctx, "黑名单规则已存在，开始删除rule: set=%s", setName)
			if delErr := nfRule.DelRule(ctx, global.Svc.Nft, nfta); delErr != nil {
				logc.Debugf(ctx, "删除set %s 的引用规则时出现提示: %v（可能规则不存在，继续执行）", setName, delErr)
			} else {
				logc.Debugf(ctx, "成功删除set %s 的旧规则", setName)
			}
		}
		// 尝试删除规则（即使规则不存在也不报错）

		// 再删除 set
		err = nfSet.DeleteSet(ctx, global.Svc.Nft)
		if err != nil {
			logc.Errorf(ctx, "删除旧的nfset %s 失败，错误信息: %v", setName, err)
			return err
		}
		logc.Debugf(ctx, "删除旧的set %s 成功", setName)
	}

	logc.Infof(ctx, "清理已存在的黑名单规则和set完成")
	return nil
}

// createBlacklistRules 创建黑名单规则和set
// 该函数会为每个set创建对应的IP集合和过滤规则
func (h *HaproxyBlackList) createBlacklistRules(ctx context.Context, ipListMap map[string]types.StringList, nfta *nftables.Table, nfChain *dfxnftables.NfChain) error {
	logc.Infof(ctx, "开始创建黑名单规则和set")

	// 遍历 ipListMap 创建 nfset 和 nfchainrule
	for setName, blacklistInfo := range ipListMap {
		// 创建 nfset
		nfSet := &dfxnftables.NfSet{
			NfMeta:  *nfMeate,
			SetName: setName,
		}

		// 创建新的 set
		err := nfSet.CreateSet(ctx, global.Svc.Nft)
		if err != nil {
			logc.Errorf(ctx, "创建nfset %s 失败，错误信息: %v", setName, err)
			return err
		}
		logc.Debugf(ctx, "创建set %s 成功", setName)

		// 添加 IP 地址到 set
		if len(blacklistInfo.Values) > 0 {
			// 分批插入，每批最多100个IP
			batchSize := 100
			totalIPs := len(blacklistInfo.Values)

			for i := 0; i < totalIPs; i += batchSize {
				end := i + batchSize
				if end > totalIPs {
					end = totalIPs
				}

				batch := blacklistInfo.Values[i:end]
				nfSet.SetIPList = batch
				err = nfSet.AddElement(ctx, global.Svc.Nft)
				if err != nil {
					logc.Errorf(ctx, "添加IP到set %s 失败 (批次 %d-%d)，错误信息: %v",
						setName, i, end, err)
					return err
				}

				logc.Debugf(ctx, "成功添加 %d-%d/%d 个IP到set %s",
					i+1, end, totalIPs, setName)
			}

			logc.Infof(ctx, "成功添加 %d 个IP到set %s", totalIPs, setName)
		}

		// 解析端口
		var dstPorts []uint16
		if blacklistInfo.Port != "" {
			var port uint16
			_, err := fmt.Sscanf(blacklistInfo.Port, "%d", &port)
			if err == nil && port > 0 {
				dstPorts = []uint16{port}
			}
		}

		nfchd, err := nfChain.QueryChainHandle(ctx, global.Svc.Nft, nfta)
		if err != nil {
			logc.Errorf(ctx, "查询nfchain %s 失败，错误信息: %v", nfMeate.TableName, err)
			return err
		}

		// 创建 nfchainrule - 匹配源IP在set中的流量并drop
		nfRule := &dfxnftables.NfChainRule{
			NfMeta:    *nfMeate,
			ChainName: "input",
			SrcSet:    setName,
			Protocol:  "tcp",
			DstPorts:  dstPorts,
			Verdict:   "drop",
		}

		nfub, err := nfRule.CheckRule(ctx, global.Svc.Nft, nfta, nfchd)
		if err == nil && nfub {
			logc.Infof(ctx, "黑名单规则已存在，跳过创建: set=%s, port=%s", setName, blacklistInfo.Port)
			continue
		}

		err = nfRule.AddRule(ctx, global.Svc.Nft, nfta)
		if err != nil {
			logc.Errorf(ctx, "创建nfchainrule for set %s 失败，错误信息: %v", setName, err)
			return err
		}

		logc.Infof(ctx, "成功创建黑名单规则: set=%s, port=%s, ip_count=%d",
			setName, blacklistInfo.Port, len(blacklistInfo.Values))
	}

	logc.Infof(ctx, "创建黑名单规则和set完成")
	return nil
}

func (h *HaproxyBlackList) EnableBlacklist(
	ctx context.Context,
) error {
	ipListMap, err := GetServerBlackListApiIpListMap(ctx, global.CMDB_BLACKLIST_HAPROXY_NAME)
	if err != nil {
		logc.Errorf(ctx, "获取服务端黑名单数据失败，错误信息: %v", err)
		return err
	}

	err = nfTable.DeleteTables(ctx, global.Svc.Nft)
	if err != nil {
		logc.Errorf(ctx, "删除nftables失败，错误信息: %v", err)
		return err
	}

	
	common.LogDebugJSON(ctx, global.Svc.Config.LogConf.Level, ipListMap)

	err = nfTable.CreateTables(ctx, global.Svc.Nft)
	if err != nil {
		logc.Errorf(ctx, "创建nftables失败，错误信息: %v", err)
		return err
	}

	nfta, err := nfTable.QueryTableshandle(ctx, global.Svc.Nft)
	if err != nil {
		logc.Errorf(ctx, "查询nftables失败，错误信息: %v", err)
		return err
	}

	// 创建链配置对象
	policy := nftables.ChainPolicyAccept
	chainPriority := *nftables.ChainPriorityFilter
	nfChain := &dfxnftables.NfChain{
		NfMeta:    *nfMeate,
		ChainName: "input",
		Type:      nftables.ChainTypeFilter,
		Hook:      nftables.ChainHookInput,
		Priority:  int32(chainPriority),
		Policy:    &policy,
	}

	err = nfChain.CreateChain(ctx, global.Svc.Nft)
	if err != nil {
		logc.Errorf(ctx, "创建nfchain失败，错误信息: %v", err)
		return err
	}

	// 第二步：创建新的规则和set
	if err := h.createBlacklistRules(ctx, ipListMap, nfta, nfChain); err != nil {
		logc.Errorf(ctx, "创建黑名单规则失败: %v", err)
		return err
	}

	return nil
}

func (h *HaproxyBlackList) DisableBlacklist(
	ctx context.Context,
) error {
	err := nfTable.DeleteTables(ctx, global.Svc.Nft)
	if err != nil {
		logc.Errorf(ctx, "删除nftables失败，错误信息: %v", err)
		return err
	}
	return nil
}
